import torch
import cv2
import os
import sys
import gc
import numpy as np
from ultralytics import YOLO
from threading import Thread
import logging

from nodes.Node import Node
from nodes.YOLONode import YOLONode
import time

from multiprocessing import Process, Queue

from nodes.SORTTrackNode import SORTTrackNode

class FishBatchYOLONode(YOLONode):

    def __init__(self,name="FishBatchYOLONode",parent=None):
        super().__init__(name, parent)
        self._batchSize = 10

    def setBatchSize(self, size):
        self._batchSize = size

    def run(self):
        self.setStatus(Node.NODE_STATUS_RUNNING)
        device = torch.device("cuda:0")
        torch.cuda.set_device(device)
        self.open()

        self.is_stopped = False
        while True:
            if self.is_stopped:
                if self.is_opened:
                    self.close()
                break
            else:
                frames = []
                values = []
                results = []

                bit = 0
                while bit < self._batchSize:
                    if( not self._framePool.empty() ):
                        frame, *value = self._framePool.get()
                        frames.append(frame)
                        values.append(value)
                        bit+=1
                    else:
                        gc.collect()
                        if(self.getParentStatus() == Node.NODE_STATUS_EXITED):
                            break

                if len(frames) < 1:
                    continue

                #print("Len:",len(frames))

                results = self._model(frames,imgsz=640,device=0)
                #results = results.boxes.cpu()
                #results = results.cpu().numpy()

                for index, frame in enumerate(frames):
                    if frame is not None:
                        videoIndex = values[index][0]
                        frameIndex = values[index][1]
                        self.passFrame([frame, videoIndex, frameIndex, results[index].boxes.cpu().numpy()])

                if self._framePool.empty():
                    if self.getParentStatus() == Node.NODE_STATUS_EXITED:
                        logging.error("SourceNode is already finished.")
                        break

                del frames
                del values
                del results

        self.close()
        self.setStatus(Node.NODE_STATUS_EXITED)
        logging.error("++++++++++++++++ FishBatchYOLONode is completed. +++++++++++++++")
        #sys.exit()
